import torch
import math
import random
from torchvision.transforms.v2 import GaussianBlur


class CustomGaussianBlur:
    def __init__(self, img_size=84, min_strength=0.01, max_strength=0.05):
        """
        Initialize the CustomGaussianBlur with a range for the blur strength.
        
        Args:
        img_size (int): Size of the image (assuming square images for simplicity).
        min_strength (float): Minimum factor for determining the kernel size as a percentage of the image size.
        max_strength (float): Maximum factor for determining the kernel size as a percentage of the image size.
        """
        self.img_size = img_size
        self.min_strength = min_strength
        self.max_strength = max_strength

    def _compute_kernel_size(self, strength):
        """
        Compute the kernel size based on the image size and a given strength.

        Args:
        strength (float): Strength factor to determine the kernel size.

        Returns:
        int: Computed kernel size ensuring it's odd.
        """
        kernel_size = math.ceil(self.img_size * strength)
        kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size
        return kernel_size

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply Gaussian Blur to the input tensor with a randomly selected strength.
        
        Args:
        x (torch.Tensor): Input tensor with shape [T, H, W].

        Returns:
        torch.Tensor: Blurred tensor with the same shape as input after dynamically applying Gaussian blur.
        """
        strength = random.uniform(self.min_strength, self.max_strength)
        kernel_size = self._compute_kernel_size(strength)
        gaussian_blur = GaussianBlur(kernel_size=kernel_size)

        return gaussian_blur(x.unsqueeze(1)).squeeze(1)
